{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d76a68dd",
   "metadata": {},
   "source": [
    "# Chat Method for Specific JSON Format\n",
    "This notebook demonstrates how to call a chat method to return data in a specific JSON format. We'll use OpenAI's API to structure the response in a predefined format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8a465700",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Required Libraries\n",
    "import openai\n",
    "import json\n",
    "from jsonschema import validate\n",
    "import os\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2490d25a",
   "metadata": {},
   "source": [
    "# Define the Desired JSON Format\n",
    "We'll create a schema that defines our expected JSON structure with indicators, dimensions, and filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a1bbde19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the expected JSON schema\n",
    "json_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"indicators\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\"type\": \"string\"}\n",
    "        },\n",
    "        \"dimensions\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\"type\": \"string\"}\n",
    "        },\n",
    "        \"filters\": {\n",
    "            \"type\": \"array\",\n",
    "            \"items\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"field\": {\"type\": \"string\"},\n",
    "                    \"value\": {\"type\": \"string\"}\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5451ff2",
   "metadata": {},
   "source": [
    "# Set Up the Chat Method\n",
    "Configure the OpenAI chat method with a prompt that instructs the model to return the JSON format we want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "11bbe90d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load environment variables\n",
    "load_dotenv()\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "def create_chat_prompt(query):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": \"\"\"You are a data analysis assistant. Return responses in the following JSON format:\n",
    "        {\n",
    "            \"indicators\": [\"list of metrics\"],\n",
    "            \"dimensions\": [\"list of dimensions\"],\n",
    "            \"filters\": [{\"field\": \"filter_name\", \"value\": \"filter_value\"}]\n",
    "        }\n",
    "        \"\"\"},\n",
    "        {\"role\": \"user\", \"content\": query}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "778912c6",
   "metadata": {},
   "source": [
    "# Call the Chat Method\n",
    "Test the implementation with a sample query and capture the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7ddd68ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: \n",
      "\n",
      "You tried to access openai.ChatCompletion, but this is no longer supported in openai>=1.0.0 - see the README at https://github.com/openai/openai-python for the API.\n",
      "\n",
      "You can run `openai migrate` to automatically upgrade your codebase to use the 1.0.0 interface. \n",
      "\n",
      "Alternatively, you can pin your installation to the old version, e.g. `pip install openai==0.28`\n",
      "\n",
      "A detailed migration guide is available here: https://github.com/openai/openai-python/discussions/742\n",
      "\n",
      "null\n"
     ]
    }
   ],
   "source": [
    "def get_structured_response(query):\n",
    "    try:\n",
    "        response = openai.ChatCompletion.create(\n",
    "            model=\"gpt-3.5-turbo\",\n",
    "            messages=create_chat_prompt(query),\n",
    "            temperature=0\n",
    "        )\n",
    "        return json.loads(response.choices[0].message.content)\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "# Test with a sample query\n",
    "sample_query = \"Show me sales by region for last year\"\n",
    "result = get_structured_response(sample_query)\n",
    "print(json.dumps(result, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4ce321c",
   "metadata": {},
   "source": [
    "# Validate the Output\n",
    "Ensure the response matches our expected JSON schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c9af6a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_response(response):\n",
    "    try:\n",
    "        validate(instance=response, schema=json_schema)\n",
    "        print(\"Response validation successful!\")\n",
    "        return True\n",
    "    except Exception as e:\n",
    "        print(f\"Validation error: {str(e)}\")\n",
    "        return False\n",
    "\n",
    "# Validate the test response\n",
    "if result:\n",
    "    validate_response(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chat-bi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
